package net.osin.crypto

import java.io.ByteArrayOutputStream
import java.io.File
import java.io.FileInputStream
import java.io.FileOutputStream
import java.io.IOException
import java.io.ObjectInputStream
import java.io.ObjectOutputStream
import java.io.Serializable
import java.io.UnsupportedEncodingException
import java.math.BigInteger
import java.security.SecureRandom
import java.util

import org.bouncycastle.crypto.params.ECDomainParameters
import org.bouncycastle.math.ec.ECCurve
import org.bouncycastle.math.ec.ECPoint

import scala.collection.mutable.ArrayBuffer


/**
 * SM2公钥加密算法实现 包括 -签名,验签 -密钥交换 -公钥加密,私钥解密
 *
 * @author Potato
 *
 */
class SM2() {
  private val n = BigInt("FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "7203DF6B" + "21C6052B" + "53BBF409" + "39D54123", 16)
  private val p = BigInt("FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "00000000" + "FFFFFFFF" + "FFFFFFFF", 16)
  private val a = BigInt("FFFFFFFE" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "FFFFFFFF" + "00000000" + "FFFFFFFF" + "FFFFFFFC", 16)
  private val b = BigInt("28E9FA9E" + "9D9F5E34" + "4D5A9E4B" + "CF6509A7" + "F39789F5" + "15AB8F92" + "DDBCBD41" + "4D940E93", 16)
  private val gx = BigInt("32C4AE2C" + "1F198119" + "5F990446" + "6A39C994" + "8FE30BBF" + "F2660BE1" + "715A4589" + "334C74C7", 16)
  private val gy = BigInt("BC3736A2" + "F4F6779C" + "59BDCEE3" + "6B692153" + "D0A9877C" + "C62A4740" + "02DF32E5" + "2139F0A0", 16)
  private var ecc_bc_spec : ECDomainParameters= null
  private val w :Int= Math.ceil(n.bitLength * 1.0 / 2).toInt - 1
  private val _2w = BigInt("2").pow(w)
  private val DIGEST_LENGTH = 32
  private val random = new SecureRandom
  private var curve : ECCurve.Fp= _
  private var G : ECPoint= _
  private var debug = false


    curve = new ECCurve.Fp(p, // q
      a, // a
      b) // b
    G = curve.createPoint(gx, gy)
    ecc_bc_spec = new ECDomainParameters(curve, G, n)


  def this(debug: Boolean) {
    this
    this.debug = debug
  }

  def isDebug: Boolean = debug

  def setDebug(debug: Boolean): Unit = {
    this.debug = debug
  }
  implicit def bigInt2bigInter(b:BigInt) :BigInteger = b.bigInteger
  /**
   * 以16进制打印字节数组
   *
   * @param b
   */
  def printHexString(b: Array[Byte]): Unit = b.map("%02x".format(_)).mkString


  /**
   * 随机数生成器
   *
   * @param max
   * @return
   */
  private def random(max: BigInt) = {
    var r = BigInt(256, scala.util.Random)
    // int count = 1;
    while ( {
      r.compareTo(max) >= 0
    }) {
      r = BigInt(128, scala.util.Random)
      // count++;
    }
    // System.out.println("count: " + count);
    r
  }

  /**
   * 判断字节数组是否全0
   *
   * @param buffer
   * @return
   */
  private def allZero(buffer: Array[Byte]): Boolean = {
    for (i <- 0 until buffer.length) {
      if (buffer(i) != 0) return false
    }
    true
  }



  /**
   * 公钥加密
   *
   * @param input
   * 加密原文
   * @param publicKey
   * 公钥
   * @return
   */
  def encrypt(input: String, publicKey: ECPoint): Array[Byte] = {
    val inputBuffer = input.getBytes
    if (debug) printHexString(inputBuffer)
    var C1Buffer :Array[Byte]= null
    var kpb :ECPoint= null
    var t :Array[Byte]= null
    do {
      /* 1 产生随机数k，k属于[1, n-1] */
      val k = random(n)
      if (debug) {
        System.out.print("k: ")
        printHexString(k.toByteArray)
      }
      /* 2 计算椭圆曲线点C1 = [k]G = (x1, y1) */

      val C1 = G.multiply(k)
      C1Buffer = C1.getEncoded(false)
      if (debug) {
        System.out.print("C1: ")
        printHexString(C1Buffer)
      }

      /*
      * 3 计算椭圆曲线点 S = [h]Pb
      */
      val h = ecc_bc_spec.getH
      if (h != null) {
        val S = publicKey.multiply(h)
        if (S.isInfinity) throw new IllegalStateException
      }

      /* 4 计算 [k]PB = (x2, y2) */
      kpb = publicKey.multiply(k).normalize

      /* 5 计算 t = KDF(x2||y2, klen) */
      val kpbBytes = kpb.getEncoded(false)
      t = KDF(kpbBytes, inputBuffer.length)
      // DerivationFunction kdf = new KDF1BytesGenerator(new
      // ShortenedDigest(new SHA256Digest(), DIGEST_LENGTH));
      //
      // t = new byte[inputBuffer.length];
      // kdf.init(new ISO18033KDFParameters(kpbBytes));
      // kdf.generateBytes(t, 0, t.length);
    } while (allZero(t))

    /* 6 计算C2=M^t */
    val C2 = new Array[Byte](inputBuffer.length)
    for (i <- 0 until inputBuffer.length) {
      C2(i) = (inputBuffer(i) ^ t(i)).toByte
    }

    /* 7 计算C3 = Hash(x2 || M || y2) */
    val C3 = sm3hash(Exchange.join(kpb.getXCoord.toBigInteger.toByteArray, inputBuffer, kpb.getYCoord.toBigInteger.toByteArray))

    /* 8 输出密文 C=C1 || C2 || C3 */
    val encryptResult = new Array[Byte](C1Buffer.length + C2.length + C3.length)
    System.arraycopy(C1Buffer, 0, encryptResult, 0, C1Buffer.length)
    System.arraycopy(C2, 0, encryptResult, C1Buffer.length, C2.length)
    System.arraycopy(C3, 0, encryptResult, C1Buffer.length + C2.length, C3.length)
    if (debug) {
      System.out.print("密文: ")
      printHexString(encryptResult)
    }
    encryptResult
  }

  /**
   * 私钥解密
   *
   * @param encryptData
   * 密文数据字节数组
   * @param privateKey
   * 解密私钥
   * @return
   */
  def decrypt(encryptData: Array[Byte], privateKey: BigInt): String = {
    if (debug) System.out.println("encryptData length: " + encryptData.length)
    val C1Byte = new Array[Byte](65)
    System.arraycopy(encryptData, 0, C1Byte, 0, C1Byte.length)
    val C1 = curve.decodePoint(C1Byte).normalize
    /*
    * 计算椭圆曲线点 S = [h]C1 是否为无穷点
    */
    val h = ecc_bc_spec.getH
    if (h != null) {
      val S = C1.multiply(h)
      if (S.isInfinity) throw new IllegalStateException
    }
    /* 计算[dB]C1 = (x2, y2) */
    val dBC1 = C1.multiply(privateKey).normalize
    /* 计算t = KDF(x2 || y2, klen) */
    val dBC1Bytes = dBC1.getEncoded(false)
    val klen = encryptData.length - 65 - DIGEST_LENGTH
    val t = KDF(dBC1Bytes, klen)

    if (allZero(t)) {
      System.err.println("all zero")
      throw new IllegalStateException
    }

    /* 5 计算M'=C2^t */
    val M = new Array[Byte](klen)
    for (i <- 0 until M.length) {
      M(i) = (encryptData(C1Byte.length + i) ^ t(i)).toByte
    }
    if (debug) printHexString(M)

    /**
     *  6 计算 u = Hash(x2 || M' || y2) 判断 u == C3是否成立
     */
    val C3 = new Array[Byte](DIGEST_LENGTH)
    if (debug) try {
      System.out.println("M = " + new String(M, "UTF8"))
    } catch {
      case e1: UnsupportedEncodingException => e1.printStackTrace()
    }
    System.arraycopy(encryptData, encryptData.length - DIGEST_LENGTH, C3, 0, DIGEST_LENGTH)
    val u = sm3hash(Exchange.join(dBC1.getXCoord.toBigInteger.toByteArray, M, dBC1.getYCoord.toBigInteger.toByteArray))
    if (util.Arrays.equals(u, C3)) {
      if (debug) System.out.println("解密成功")
      try return new String(M, "UTF8")
      catch {
        case e: UnsupportedEncodingException =>
          e.printStackTrace()
      }
      null
    }
    else {
      if (debug) {
        System.out.print("u = ")
        printHexString(u)
        System.out.print("C3 = ")
        printHexString(C3)
        System.err.println("解密验证失败")
      }
      null
    }
  }


  /**
   * 判断是否在范围内
   *
   * @param param
   * @param min
   * @param max
   * @return
   */
  private def between(param: BigInt, min: BigInt, max: BigInt) = {
    if (param.compareTo(min) >= 0 && param.compareTo(max) < 0) true
    else false
  }

    /**
   * 判断生成的公钥是否合法
   *
   * @param publicKey
   * @return
   */
  private def checkPublicKey(publicKey: ECPoint): Boolean = {
    if (!publicKey.isInfinity) {
      val x = publicKey.getXCoord.toBigInteger
      val y = publicKey.getYCoord.toBigInteger
      if (between(x, BigInt("0"), p) && between(y, BigInt("0"), p)) {
        val xResult = x.pow(3).add(a*x).add(b).mod(p)
        if (debug) System.out.println("xResult: " + xResult.toString)
        val yResult = y.pow(2).mod(p)
        if (debug) System.out.println("yResult: " + yResult.toString)
        if (yResult == xResult && publicKey.multiply(n).isInfinity) return true
      }
    }
    false
  }

  /**
   * 生成密钥对
   *
   * @return
   */
  def generateKeyPair: SM2KeyPair = {
    val d = random(n-(BigInt("1")))
    val keyPair = new SM2KeyPair(G.multiply(d).normalize, d)
    if (checkPublicKey(keyPair.getPublicKey)) {
      if (debug) System.out.println("generate key successfully")
      keyPair
    }
    else {
      if (debug) System.err.println("generate key failed")
      null
    }
  }


  /**
   * 导出公钥到本地
   *
   * @param publicKey
   * @param path
   */
  def exportPublicKey(publicKey: ECPoint, path: String): Unit = {
    val file = new File(path)
    try {
      if (!file.exists) file.createNewFile
      val buffer = publicKey.getEncoded(false)
      val fos = new FileOutputStream(file)
      fos.write(buffer)
      fos.close()
    } catch {
      case e: IOException =>
        e.printStackTrace()
    }
  }

  /**
   * 从本地导入公钥
   *
   * @param path
   * @return
   */
  def importPublicKey(path: String): ECPoint = {
    val file = new File(path)
    try {
      if (!file.exists) return null
      val fis = new FileInputStream(file)
      val baos = new ByteArrayOutputStream
      val buffer = new Array[Byte](16)
      var size = 0
      while ((size = fis.read(buffer)) != -1) baos.write(buffer, 0, size)
      fis.close()
      return curve.decodePoint(baos.toByteArray)
    } catch {
      case e: IOException =>
        e.printStackTrace()
    }
    null
  }

  /**
   * 导出私钥到本地
   *
   * @param privateKey
   * @param path
   */
  def exportPrivateKey(privateKey: BigInt, path: String): Unit = {
    val file = new File(path)
    try {
      if (!file.exists) file.createNewFile
      val oos = new ObjectOutputStream(new FileOutputStream(file))
      oos.writeObject(privateKey)
      oos.close()
    } catch {
      case e: IOException =>
        e.printStackTrace()
    }
  }

  /**
   * 从本地导入私钥
   *
   * @param path
   * @return
   */
  def importPrivateKey(path: String): BigInt = {
    val file = new File(path)
    try {
      if (!file.exists) return null
      val fis = new FileInputStream(file)
      val ois = new ObjectInputStream(fis)
      val res = ois.readObject.asInstanceOf[BigInt]
      ois.close()
      fis.close()
      return res
    } catch {
      case e: Exception =>
        e.printStackTrace()
    }
    null
  }


  /**
   * sm3摘要
   *
   * @param params
   * @return
   */
  private def sm3hash(params: Array[Byte]) = SM3.hash(params)


  /**
   * 取得用户标识字节数组
   *
   * @param IDA
   * @param aPublicKey
   * @return
   */
  private def getZA(IDA: String, aPublicKey: ECPoint) = {
    val idaBytes = IDA.getBytes
    val entlenA = idaBytes.length * 8
    val ENTLA = Array[Byte]((entlenA & 0xFF00).toByte, (entlenA & 0x00FF).toByte)
    sm3hash(Exchange.join(ENTLA, idaBytes, a.toByteArray, b.toByteArray, gx.toByteArray, gy.toByteArray, aPublicKey.getXCoord.toBigInteger.toByteArray, aPublicKey.getYCoord.toBigInteger.toByteArray))
  }


  /**
   * 签名
   *
   * @param M
   * 签名信息
   * @param IDA
   * 签名方唯一标识
   * @param keyPair
   * 签名方密钥对
   * @return 签名
   */
  def sign(M: String, IDA: String, keyPair: SM2KeyPair): Signature = {
    val ZA :Array[Byte]= getZA(IDA,keyPair.getPublicKey)
    val M_ = Exchange.join(ZA, M.getBytes)
    val e = BigInt(1, sm3hash(M_))
    // BigInteger k = new BigInteger(
    // "6CB28D99 385C175C 94F94E93 4817663F C176D925 DD72B727 260DBAAE
    // 1FB2F96F".replace(" ", ""), 16);
    var k :BigInt= null
    var r :BigInt= null
    do {
      k = random(n)
      val p1 = G.multiply(k).normalize
      val x1 = p1.getXCoord.toBigInteger
      r = e+(x1)
      r = r.mod(n)
    } while (r.equals( BigInt(0)) || (r+k).equals(n))
    val s = (keyPair.getPrivateKey.+(BigInt(1)).modInverse(n))
      .*((k.-(r.*(keyPair.getPrivateKey))).mod(n)).mod(n)
    new Signature(r, s)
  }

  /**
   * 验签
   *
   * @param M
   * 签名信息
   * @param signature
   * 签名
   * @param IDA
   * 签名方唯一标识
   * @param aPublicKey
   * 签名方公钥
   * @return true or false
   */
  def verify(M: String, signature: Signature, IDA: String, aPublicKey: ECPoint): Boolean = {
    if (!between(signature.r, BigInt(1), n)) return false
    if (!between(signature.s, BigInt(1), n)) return false
    val M_ = Exchange.join(getZA(IDA, aPublicKey), M.getBytes)
    val e = BigInt(1, sm3hash(M_))
    val t = signature.r.+(signature.s).mod(n)
    if (t == BigInt(0)) return false
    val p1 = G.multiply(signature.s).normalize
    val p2 = aPublicKey.multiply(t).normalize
    val x1 = p1.add(p2).normalize.getXCoord.toBigInteger
    val R = e.+(x1).mod(n)
    if (R == signature.r) return true
    false
  }


  /**
   * 密钥派生函数
   *
   * @param Z
   * @param klen
   * 生成klen字节数长度的密钥
   * @return
   */
  private def KDF(Z: Array[Byte], klen: Int): Array[Byte] = {
    var ct = 1
    val end = Math.ceil(klen * 1.0 / 32).toInt
    val baos = new ArrayBuffer[Byte]()
    try {
      for (i <- 1 until end) {
        baos.appendAll(sm3hash(Exchange.join(Z, Exchange.toByteArray(ct))))
        ct += 1
      }
      val last = sm3hash(Exchange.join(Z, Exchange.toByteArray(ct)))
      if (klen % 32 == 0) baos.appendAll(last)
      else baos.appendAll(last++Array(0.toByte, (klen % 32).toByte))
      return baos.toArray
    } catch {
      case e: Exception =>
        e.printStackTrace()
    }
    null
  }

  /**
   * 传输实体类
   *
   * @author Potato
   *
   */
class TransportEntity(r: Array[Byte], s: Array[Byte], z: Array[Byte], pKey: ECPoint) extends Serializable {
    val R: Array[Byte] =r//R点
    val S: Array[Byte] =s //验证S
    val Z: Array[Byte] =z //用户标识
    val K: Array[Byte] = pKey.getEncoded(false)

  }

  /**
   * 密钥协商辅助类
   *
   * @author Potato
   *
   */
  class KeyExchange {

    private[crypto] var rA :BigInt= null
    private[crypto] var RA :ECPoint= null
    private[crypto] var V :ECPoint= null
    private[crypto] var Z :Array[Byte]= null
    private[crypto] var key :Array[Byte]= null
    var ID: String =_
    var keyPair: SM2KeyPair =_

    def this(ID: String, keyPair: SM2KeyPair){
      this
      this.ID=ID;
    this.keyPair = keyPair
    this.Z=getZA(ID, keyPair.getPublicKey)
    }

    /**
     * 密钥协商发起第一步
     *
     * @return
     */
    def keyExchange_1(): TransportEntity = {
      rA = random(n)
      // rA=new BigInteger("83A2C9C8 B96E5AF7 0BD480B4 72409A9A 327257F1
      // EBB73F5B 073354B2 48668563".replace(" ", ""),16);
      RA = (G.multiply(rA)).normalize
      new TransportEntity(RA.getEncoded(false), null, Z, keyPair.getPublicKey)
    }

    /**
     * 密钥协商响应方
     *
     * @param entity 传输实体
     * @return
     */
    def keyExchange_2(entity: TransportEntity): TransportEntity = {
      val rB = random(n)
      // BigInteger rB=new BigInteger("33FE2194 0342161C 55619C4A 0C060293
      // D543C80A F19748CE 176D8347 7DE71C80".replace(" ", ""),16);
      val RB = G.multiply(rB).normalize
      this.rA = rB
      this.RA = RB
      var x2 = RB.getXCoord.toBigInteger
      x2 = _2w.+(x2.and(_2w.-(BigInt(1))))
      val tB = keyPair.getPrivateKey.+(x2.multiply(rB)).mod(n)
      val RA = curve.decodePoint(entity.R).normalize
      var x1 = RA.getXCoord.toBigInteger
      x1 = _2w.+(x1.and(_2w.-(BigInt(1))))
      val aPublicKey = curve.decodePoint(entity.K).normalize
      val temp = aPublicKey.add(RA.multiply(x1).normalize).normalize
      val V = temp.multiply(ecc_bc_spec.getH.multiply(tB)).normalize
      if (V.isInfinity) throw new IllegalStateException
      this.V = V
      val xV = V.getXCoord.toBigInteger.toByteArray
      val yV = V.getYCoord.toBigInteger.toByteArray
      val KB = KDF(Exchange.join(xV, yV, entity.Z, this.Z), 16)
      key = KB
      System.out.print("协商得B密钥:")
      printHexString(KB)
      val sB = sm3hash(Exchange.join(Array[Byte](0x02), yV,
        sm3hash(Exchange.join(xV, entity.Z, this.Z, RA.getXCoord.toBigInteger.toByteArray,
          RA.getYCoord.toBigInteger.toByteArray, RB.getXCoord.toBigInteger.toByteArray,
          RB.getYCoord.toBigInteger.toByteArray))))
      new TransportEntity(RB.getEncoded(false), sB, this.Z, keyPair.getPublicKey)
    }

    /**
     * 密钥协商发起方第二步
     *
     * @param entity 传输实体
     */
    def keyExchange_3(entity: TransportEntity): TransportEntity = {
      var x1 = RA.getXCoord.toBigInteger
      x1 = _2w.+(x1.and(_2w.-(BigInt(1))))
      val tA = keyPair.getPrivateKey.+(x1.multiply(rA)).mod(n)
      val RB = curve.decodePoint(entity.R).normalize
      var x2 = RB.getXCoord.toBigInteger
      x2 = _2w.+(x2.and(_2w.-(BigInt(1))))
      val bPublicKey = curve.decodePoint(entity.K).normalize
      val temp = bPublicKey.add(RB.multiply(x2).normalize).normalize
      val U = temp.multiply(ecc_bc_spec.getH.multiply(tA)).normalize
      if (U.isInfinity) throw new IllegalStateException
      this.V = U
      val xU = U.getXCoord.toBigInteger.toByteArray
      val yU = U.getYCoord.toBigInteger.toByteArray
      val KA = KDF(Exchange.join(xU, yU, this.Z, entity.Z), 16)
      key = KA
      System.out.print("协商得A密钥:")
      printHexString(KA)
      val s1 = sm3hash(Exchange.join(Array[Byte](0x02), yU,
        sm3hash(Exchange.join(xU, this.Z, entity.Z, RA.getXCoord.toBigInteger.toByteArray,
          RA.getYCoord.toBigInteger.toByteArray, RB.getXCoord.toBigInteger.toByteArray,
          RB.getYCoord.toBigInteger.toByteArray))))
      if (util.Arrays.equals(entity.S, s1)) System.out.println("B->A 密钥确认成功")
      else System.out.println("B->A 密钥确认失败")
      val sA = sm3hash(Exchange.join(Array[Byte](0x03), yU,
        sm3hash(Exchange.join(xU, this.Z, entity.Z, RA.getXCoord.toBigInteger.toByteArray,
          RA.getYCoord.toBigInteger.toByteArray, RB.getXCoord.toBigInteger.toByteArray,
          RB.getYCoord.toBigInteger.toByteArray))))
      new TransportEntity(RA.getEncoded(false), sA, this.Z, keyPair.getPublicKey)
    }

    /**
     * 密钥确认最后一步
     *
     * @param entity 传输实体
     */
    def keyExchange_4(entity: TransportEntity): Unit = {
      val xV = V.getXCoord.toBigInteger.toByteArray
      val yV = V.getYCoord.toBigInteger.toByteArray
      val RA = curve.decodePoint(entity.R).normalize
      val s2 = sm3hash(Exchange.join(Array[Byte](0x03), yV,
        sm3hash(Exchange.join(xV, entity.Z, this.Z, RA.getXCoord.toBigInteger.toByteArray,
          RA.getYCoord.toBigInteger.toByteArray, this.RA.getXCoord.toBigInteger.toByteArray,
          this.RA.getYCoord.toBigInteger.toByteArray))))
      if (util.Arrays.equals(entity.S, s2)) System.out.println("A->B 密钥确认成功")
      else System.out.println("A->B 密钥确认失败")
    }
  }



  class Signature(var r: BigInt, var s: BigInt) {
    override def toString: String = r.toString(16) + "," + s.toString(16)
  }


}
